import pygame
import math
import numpy as np
import sys

# ========== 参数 ==========
MAP_WIDTH, MAP_HEIGHT = 800, 600
OBSTACLES = [
    pygame.Rect(300, 200, 150, 100),
    pygame.Rect(550, 400, 120, 80),
    pygame.Rect(100, 400, 100, 60)
]
NUM_RAYS = 36
RAY_LENGTH = 200
STEP = 2
ROBOT_RADIUS = 15
WHEEL_BASE = 40

# ========== 差速机器人类 ==========
class DiffDriveRobot:
    def __init__(self, x, y, theta=0.0):
        self.x = x
        self.y = y
        self.theta = theta
        self.v_l = 0.0
        self.v_r = 0.0

    def update(self, dt=0.05):
        v = (self.v_l + self.v_r) / 2.0
        omega = (self.v_r - self.v_l) / WHEEL_BASE
        self.x += v * math.cos(self.theta) * dt
        self.y += v * math.sin(self.theta) * dt
        self.x = np.clip(self.x, 0, MAP_WIDTH)
        self.y = np.clip(self.y, 0, MAP_HEIGHT)
        self.theta += omega * dt

# ========== 环境类 ==========
class Env:
    def __init__(self):
        self.robot = DiffDriveRobot(200, 150, math.pi / 4)
        self.obstacles = OBSTACLES
        self.surface = pygame.display.set_mode((MAP_WIDTH, MAP_HEIGHT))
        pygame.display.set_caption("Top-Down Raycasting Robot Simulation")

    def ray_cast(self):
        """进行 ray casting，并返回归一化距离"""
        hits = []
        for i in range(NUM_RAYS):
            angle = self.robot.theta + (i / NUM_RAYS) * 2 * math.pi
            hit_distance = RAY_LENGTH
            for d in range(0, RAY_LENGTH, STEP):
                x = self.robot.x + d * math.cos(angle)
                y = self.robot.y + d * math.sin(angle)
                if x < 0 or x > MAP_WIDTH or y < 0 or y > MAP_HEIGHT:
                    hit_distance = d
                    break
                if any(rect.collidepoint(x, y) for rect in self.obstacles):
                    hit_distance = d
                    break
            hits.append(hit_distance)
        norm = np.array(hits) / RAY_LENGTH
        return hits, norm

    def render(self, hits):
        """渲染地图、小车、障碍物和激光"""
        self.surface.fill((255, 255, 255))

        # 绘制障碍物
        for rect in self.obstacles:
            pygame.draw.rect(self.surface, (80, 80, 80), rect)

        # 绘制射线
        for i, d in enumerate(hits):
            angle = self.robot.theta + (i / NUM_RAYS) * 2 * math.pi
            x_end = self.robot.x + d * math.cos(angle)
            y_end = self.robot.y + d * math.sin(angle)
            pygame.draw.line(self.surface, (255, 0, 0), (self.robot.x, self.robot.y), (x_end, y_end), 1)

        # 绘制机器人主体
        pygame.draw.circle(self.surface, (0, 150, 255), (int(self.robot.x), int(self.robot.y)), ROBOT_RADIUS)
        head_x = self.robot.x + ROBOT_RADIUS * math.cos(self.robot.theta)
        head_y = self.robot.y + ROBOT_RADIUS * math.sin(self.robot.theta)
        pygame.draw.line(self.surface, (30, 30, 30), (self.robot.x, self.robot.y), (head_x, head_y), 2)

        pygame.display.flip()

# ========== 主循环 ==========
def run():
    pygame.init()
    clock = pygame.time.Clock()
    env = Env()

    print("控制说明：")
    print("↑ ↓ ：前进/后退")
    print("← → ：左转/右转（通过改变左右轮速度差实现）")
    print("ESC ：退出")

    running = True
    while running:
        dt = clock.tick(60) / 1000.0
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

        # 控制输入
        keys = pygame.key.get_pressed()
        base_speed = 60.0
        turn_speed = 40.0

        if keys[pygame.K_UP]:
            env.robot.v_l = base_speed
            env.robot.v_r = base_speed
        elif keys[pygame.K_DOWN]:
            env.robot.v_l = -base_speed
            env.robot.v_r = -base_speed
        elif keys[pygame.K_LEFT]:
            env.robot.v_l = base_speed - turn_speed
            env.robot.v_r = base_speed + turn_speed
        elif keys[pygame.K_RIGHT]:
            env.robot.v_l = base_speed + turn_speed
            env.robot.v_r = base_speed - turn_speed
        else:
            env.robot.v_l = 0
            env.robot.v_r = 0

        # 更新机器人
        env.robot.update(dt)

        # Ray casting
        hits, norm = env.ray_cast()
        env.render(hits)

        # 打印部分归一化距离
        if pygame.time.get_ticks() % 500 < 20:
            print("Ray distances (normalized) =", np.round(norm, 2))

    pygame.quit()
    sys.exit()

if __name__ == "__main__":
    run()
